#include "monitor/Monitor.h"

#include <getopt.h>
#include <readline/history.h>
#include <readline/readline.h>


static const uint32_t dimg[] = {
    0x00000297, // auipc t0,0
    0x0002b823, // sd  zero,16(t0)
    0x0102b503, // ld  a0,16(t0)
    0x00100073, // ebreak (used as nemu_trap)
    0xdeadbeef, // some data
};

int Monitor::cmd_si() {
    int n = 1;
    std::string arg;
    if (iss >> arg) n = std::atoi(arg.data());
    this->run(n);
    return 0;
}

int Monitor::cmd_c() {
    this->run(-1);
    return 0;
}

int Monitor::cmd_e() {
    std::string arg;

    if (!std::getline(iss, arg)) return 0;

    auto [success, val] = expr->parser(arg.c_str());
    if (!success) {
        std::cout << "Bad Expression: " << arg << std::endl;
    } else {
        std::cout << arg << " = ";
        std::cout << "[" << std::dec << val << ", " << std::showbase << std::hex << val << "]" << std::endl;
    }
    return 0;
}

int Monitor::cmd_x() {
    std::string arg;
    if (!(iss >> arg)) return 0;

    int N = std::atoi(arg.data());

    if (!(iss >> arg)) return 0;
    auto [success, base_addr] = expr->parser(arg.c_str());
    if (!success) {
        std::cout << "Bad Expression: " << arg << std::endl;
        return 0;
    }
    int word_size = sizeof(word_t);
    for (int i = 0; i < N; ++i) {
        word_t data;
        bool mmio;
        ram->read(base_addr + i * word_size, word_size, (uint8_t*)&data, mmio);
        printf(FMT_WORD ": ", (base_addr + i * word_size));
        for (int j = 0; j < word_size; ++j) {
            printf("0x%02x  ", (uint8_t) (data&0xFF));
            data >>= word_size;
        }
        printf("\n");
    }
    return 0;
}

int Monitor::cmd_i() {
#ifdef CONFIG_ITRACE
    std::string arg;
    iss >> arg;
    if (arg == "on") {
        core.itrace->enable = true;
        return 0;
    }
    if (arg == "off") {
        core.itrace->enable = false;
        return 0;
    }
    core.itrace->print();
#endif
    return 0;
}

int Monitor::cmd_r() {
    printf(ANSI_FMT("============================= <Register> =============================\n", ANSI_FG_MAGENTA));
    printf("General Register: \n");
    for (int i = 0; i < 32; ++i) {
        printf("\t%-3s: " FMT_WORD "\t", core.gpr_name_table[i].data(), core.gpr[i]);
        if ((i+1)%2 == 0) printf("\n");
    }

    printf("Control and Status Register: \n");
    printf("\t%-3s: " FMT_WORD "\t\n", "mstatus", core.csr[MSTATUS]);
    printf("\t%-3s: " FMT_WORD "\t\n", "mie", core.csr[MIE]);
    printf("\t%-3s: " FMT_WORD "\t\n", "mtvec", core.csr[MTVEC]);
    printf("\t%-3s: " FMT_WORD "\t\n", "mepc", core.csr[MEPC]);
    printf("\t%-3s: " FMT_WORD "\t\n", "mcause", core.csr[MCAUSE]);
    printf("\t%-3s: " FMT_WORD "\t\n", "mip", core.csr[MIP]);

    printf("Program Counter Register: \n");
    printf("\t%-3s: " FMT_WORD "\t\n", "pc", core.pc);
    printf(ANSI_FMT("============================= <Register> =============================\n\n", ANSI_FG_MAGENTA));
    return 0;
}

int Monitor::cmd_q() {
    state = State::QUIT;
    return -1;
}

void Monitor::cmd_table_init() {
    expr = new Expression([&](std::string& name, bool& success) -> word_t {
        if (name == "$pc") return core.pc;
        for (int i = 0; i < 32; ++i) {
            if (name == core.gpr_name_table[i]) 
            return core.gpr[i];
        }
        success = false;
        return 0;
    });

    cmd_table.push_back({"si", "step n instructions", std::bind(&Monitor::cmd_si, this)});
    cmd_table.push_back({"c", "continue execution", std::bind(&Monitor::cmd_c, this)});
    cmd_table.push_back({"e", "evaluate expression", std::bind(&Monitor::cmd_e, this)});
    cmd_table.push_back({"x", "examine memory", std::bind(&Monitor::cmd_x, this)});
    cmd_table.push_back({"i", "print instruction trace", std::bind(&Monitor::cmd_i, this)});
    cmd_table.push_back({"r", "print register information", std::bind(&Monitor::cmd_r, this)});
    cmd_table.push_back({"q", "quit", std::bind(&Monitor::cmd_q, this)});
}

void Monitor::device_init(const char* img_file) {

    if (img_file != NULL && std::filesystem::exists(img_file)) {
        Log("The image is %s", img_file);
        ram = new Memory(img_file);
    } else {
        Log("No image is given. Use the default build-in image.\n");
        ram = new Memory((uint8_t*)dimg, sizeof(dimg));
    }
    bus.add_device(MEM_BASE, ram->get_size(), ram);

    serial = new Serial();
    bus.add_device(SERIAL_BASE, SERIAL_SIZE, serial);

    timer = new Timer();
    bus.add_device(TIMER_BASE, TIMER_SIZE, timer);

#if CONFIG_CLINT
    clint = new CLINT();
    bus.add_device(CLINT_BASE, CLINT_SIZE, clint);
#endif

#if CONFIG_KEYBOARD
    keyboard = new Keyboard();
    bus.add_device(KEYBOARD_BASE, KEYBOARD_SIZE, keyboard);
#endif

#if CONFIG_DISPLAY
    display = new Display();
    bus.add_device(VCTL_BASE, VCTL_SIZE, display);
    bus.add_device(VMEM_BASE, VMEM_SIZE, display);
#endif

    t = new std::thread([&]() -> void {
        while (state == State::RUN) {
            bus.update();
            std::this_thread::sleep_for(std::chrono::milliseconds(1000 / 60));
        }
    });
}

Monitor::Monitor(int argc, char *argv[]) : core(this->bus)
{
    char* img_file = NULL;
    char* ref_file = NULL;
    char* log_path = NULL;
    bool is_batch_mode = false;

    const struct option table[] = {
        {"batch"    , no_argument      , NULL, 'b'},
        {"img"      , required_argument, NULL, 'i'},
        {"ref"      , required_argument, NULL, 'r'},
        {"log"      , required_argument, NULL, 'l'},
        {"help"     , no_argument      , NULL, 'h'},
        {0          , 0                , NULL,  0 },
    };
    int o;
    while ( (o = getopt_long(argc, argv, "-bhi:r:", table, NULL)) != -1) {
        switch (o) {
            case 'b': is_batch_mode = true; break;
            case 'i': img_file = optarg; break;
            case 'r': ref_file = optarg; break;
            case 'l': log_path = optarg; break;
            default:
                printf("Usage: %s [OPTION...] IMAGE [args]\n\n", argv[0]);
                printf("\t-b,--batch              run with batch mode\n");
                printf("\t-l,--log=FILE           output log to FILE\n");
                printf("\n");
                exit(0);
        }
    }

    this->batch_mode = is_batch_mode;

    device_init(img_file);

    cmd_table_init();

    state = State::RUN;

    if (log_path) std::filesystem::create_directory(log_path);
    core.open_trace(log_path);

    IFDEF(CONFIG_DIFFTEST, core.difftest_init(ref_file, ram->get_base(), ram->get_size()));
}

Monitor::~Monitor()
{
    if (t->joinable()) { t->join(); delete t; };
    if (clint != nullptr) delete clint;
    if (ram != nullptr) delete ram;
    if (expr != nullptr) delete expr;
    if (serial != nullptr) delete serial;
    if (timer != nullptr) delete timer;
    if (keyboard != nullptr) delete keyboard;
    if (display != nullptr) delete display;
}


void Monitor::run(size_t n) {
    if (state == State::END || state == State::ABORT) {
        std::cout << "Program execution has ended. To restart the program, exit NEMU and run again." << std::endl;
        return;
    }
    bool msi = false, mti = false;
    while (n --) {
        try {
#ifdef CONFIG_CLINT
            clint->tick();
            msi = clint->msi(), mti = clint->mti();
#endif
            core.step(msi, mti);
        } catch (RVException& e) {
            if (e.get_trap_cause() == Trap::Breakpoint) {
                state = State::END;
                halt_value = e.get_trap_val();
                if (halt_value == 0) {
                    std::printf(ANSI_FMT("HIT GOOD TRAP", ANSI_FG_GREEN) "\n");
                } else {
                    std::printf(ANSI_FMT("HIT BAD TRAP", ANSI_FG_RED) "\n");
                }
            } else {
                std::cout << e.what() << std::endl;
                state = State::ABORT;
                std::printf(ANSI_FMT("ABORT", ANSI_FG_RED) " at pc = " FMT_WORD "\n", core.pc);
            }
            return;
        }
    }
}

int Monitor::loop() {
    if (batch_mode) {
        this->run(-1);
        return !(state == State::END && halt_value == 0 || state == State::QUIT);
    }
    std::string line, cmd;
    std::cout << "Welcome to npc!" << std::endl;
    using_history();

    while (true) {
        char* s = readline("(zemu) ");
        if (s && *s) { add_history(s); line = s; }
        if (s != NULL) { free(s); }
        else { continue; }

        iss.clear();
        iss.str(line);
        iss >> cmd;

        auto cmd_it = std::find_if(cmd_table.begin(), cmd_table.end(), [&](Command& c) -> bool {
            return c.name == cmd;
        });
        if (cmd_it == cmd_table.end()) {
            std::cout << "Unknown command: " << cmd << std::endl;
            continue;
        }
        if (cmd_it->handler() < 0) { break; }
    }
    return !(state == State::END && halt_value == 0 || state == State::QUIT);
}
